{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# 邮件攻防战\n",
    "\n",
    "* 数据下载[点击这里](https://trec.nist.gov/data.html)\n",
    "* SpamBayes相关下载[点击这里](https://github.com/jfinkels/sbclassifier)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from sbclassifier import Classifier\n",
    "from sbclassifier.classifiers.constants import HAM_CUTOFF\n",
    "from sbclassifier.classifiers.constants import SPAM_CUTOFF"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 1. 建立 SpamBayes 邮件分类系统"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "h_lim = 2500\n",
    "s_lim = 5000\n",
    "train_ratio = 0.5\n",
    "vocab_ratio = train_ratio + 0.3\n",
    "symbol_list = ['_' , '-', ',' , '.' , '(' , ')', '[' , ']' , '{' , '}']"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### 1.1 建立训练语料库"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def create_train_tokens():\n",
    "    with open('data/trec05/index') as f:\n",
    "        index_file = f.readlines()\n",
    "    h_tokens = []\n",
    "    s_tokens = []\n",
    "    h_num = 0\n",
    "    s_num = 0\n",
    "    \n",
    "    # 取邮件数据\n",
    "    for index in index_file:\n",
    "        try:\n",
    "            if index.startswith('h') and h_num < h_lim:\n",
    "                h_num += 1\n",
    "                this_token = readmail('data/trec05' + index.split('..')[1].strip('\\n'))\n",
    "                h_tokens.append(this_token)\n",
    "            if index.startswith('s') and s_num < s_lim:\n",
    "                s_num += 1\n",
    "                this_token = readmail('data/trec05' + index.split('..')[1].strip('\\n'))\n",
    "                s_tokens.append(this_token)               \n",
    "            if h_num == h_lim and s_num == s_lim:\n",
    "                print(\"获得了垃圾邮件{}封\".format(s_num))\n",
    "                print(\"获得了订阅邮件{}封\".format(h_num))\n",
    "                print(\"\\n-------------我是分割线------------\\n\")\n",
    "                break\n",
    "        except:\n",
    "            continue\n",
    "            \n",
    "    # 数据划分\n",
    "    train_h = h_tokens[0 : int(h_lim*train_ratio)]\n",
    "    vocab_h = h_tokens[int(h_lim*train_ratio) : int(h_lim*vocab_ratio)]\n",
    "    test_h = h_tokens[int(h_lim*vocab_ratio) : ]\n",
    "    train_s = s_tokens[0 : int(s_lim*train_ratio)]\n",
    "    vocab_s = s_tokens[int(s_lim*train_ratio) : int(s_lim*vocab_ratio)]\n",
    "    test_s = s_tokens[int(s_lim*vocab_ratio) : ]\n",
    "    \n",
    "    print(\"训练用 ham {} 封\".format(len(train_h)))\n",
    "    print(\"攻击用 ham {} 封\".format(len(vocab_h)))\n",
    "    print(\"测试用 ham {} 封\".format(len(test_h)))\n",
    "    print(\"训练用 spam {} 封\".format(len(train_s)))\n",
    "    print(\"攻击用 spam {} 封\".format(len(vocab_s)))\n",
    "    print(\"测试用 spam {} 封\".format(len(test_s)))\n",
    "    \n",
    "    return train_h, train_s, vocab_h, vocab_s, test_h, test_s\n",
    "\n",
    "def readmail(mail_path):\n",
    "    with open(mail_path, 'r') as mail_f:\n",
    "        mail = mail_f.read()\n",
    "        for symbol in symbol_list:\n",
    "            mail = mail.replace(symbol, ' ')\n",
    "        tokens = mail.split()\n",
    "    return tokens\n",
    "\n",
    "train_h, train_s, vocab_h, vocab_s, test_h, test_s = create_train_tokens()\n",
    "print(\"\\n训练库 | 攻击库 | 测试库 已创建！\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### 1.2 完成模型训练"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def train_model(h_tokens, s_tokens):\n",
    "    model = Classifier()\n",
    "    for ham in h_tokens:\n",
    "        model.learn_ham(ham)\n",
    "    for spam in s_tokens:\n",
    "        model.learn_spam(spam)\n",
    "    return model\n",
    "\n",
    "model = train_model(train_h, train_s)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### 1.3 尝试得分预测"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "model.spamprob(['this','is','my','adv'])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "model.spamprob(test_s[44])"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### 1.4 建立评估代码"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "HAM_CUTOFF"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "SPAM_CUTOFF"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def eval(model, h_tokens, s_tokens):\n",
    "    score = 0\n",
    "    for token in h_tokens:\n",
    "        h_prob = model.spamprob(token)\n",
    "        if h_prob <= HAM_CUTOFF:\n",
    "            score += 1\n",
    "    for token in s_tokens:\n",
    "        s_prob = model.spamprob(token)\n",
    "        if s_prob >= SPAM_CUTOFF:\n",
    "            score += 1\n",
    "    result = score / (len(h_tokens) + len(s_tokens))\n",
    "    #print(\"当前模型准确率为：{:.2f}%\".format(result*100))\n",
    "    return result"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "score = eval(model, test_h, test_s)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 2. 开始攻击"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import itertools\n",
    "import numpy as np\n",
    "from matplotlib import pyplot as plt"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### 2.1 创建字典\n",
    "\n",
    "* 碾平 list 方法性能比较[点击这里](https://stackoverflow.com/questions/952914/how-to-make-a-flat-list-out-of-list-of-lists)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def plot_scores(scores):\n",
    "    plt.plot(np.linspace(0,len(scores),len(scores)),scores,'b-')\n",
    "    plt.show()\n",
    "    \n",
    "def create_vocablist(h_tokens, s_tokens):\n",
    "    h_vocab_set = list(set(itertools.chain(*vocab_h)))\n",
    "    s_vocab_set = list(set(itertools.chain(*vocab_s)))\n",
    "    print(len(h_vocab_set), len(s_vocab_set))\n",
    "    return h_vocab_set, s_vocab_set\n",
    "\n",
    "h_vocab_set, s_vocab_set = create_vocablist(vocab_h, vocab_s)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### 2.2 Dictionary Attack"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "h_vocab_arange = range(1000,2000)\n",
    "s_vocab_arange = range(1000,2000)\n",
    "atk_total = 2000"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# 初始化模型\n",
    "model = train_model(train_h, train_s)\n",
    "\n",
    "def dictionary_atk():\n",
    "    scores = []\n",
    "    print(\"攻击前模型对测试数据的识别准确率为：\")\n",
    "    old_result = eval(model, test_h, test_s)\n",
    "    print(\"\\n-------------我是分割线------------\\n\")\n",
    "    \n",
    "    for n in range(atk_total):\n",
    "        # 创建攻击邮件\n",
    "        h_num = np.random.choice(h_vocab_arange)\n",
    "        s_num = np.random.choice(s_vocab_arange)\n",
    "        h_chars = np.random.choice(h_vocab_set, h_num)\n",
    "        s_chars = np.random.choice(s_vocab_set, s_num)\n",
    "        chars = np.hstack((h_chars, s_chars))\n",
    "        np.random.shuffle(chars)\n",
    "        atk_mail = chars.tolist()\n",
    "        print(\"已创建第 {} 封攻击邮件，邮件长度为 {}\".format(n, len(atk_mail)))\n",
    "        \n",
    "        # 送入邮件防护系统进行识别\n",
    "        mail_result = model.spamprob(atk_mail)\n",
    "        \n",
    "        # 根据得分对模型进行正常的再训练\n",
    "        if mail_result <= HAM_CUTOFF:\n",
    "            print(\"邮件被识别为 ham\")\n",
    "            model.learn_ham(atk_mail)\n",
    "        elif mail_result >= SPAM_CUTOFF:\n",
    "            print(\"邮件被识别为 spam\")\n",
    "            model.learn_spam(atk_mail)\n",
    "        else:\n",
    "            print(\"邮件被识别为 unknown\")\n",
    "            \n",
    "        if n % 10 == 0:\n",
    "            _result = eval(model, test_h, test_s)\n",
    "            scores.append(_result)\n",
    "            \n",
    "    print(\"\\n-------------我是分割线------------\\n\")\n",
    "    print(\"攻击后模型对测试数据识别准确率为：\")\n",
    "    result = eval(model, test_h, test_s)\n",
    "    return scores\n",
    "\n",
    "dictionary_scores = dictionary_atk()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "plot_scores(dictionary_scores)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### 2.3 Focused Attack"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "focused_mail = np.random.choice(test_h)\n",
    "model = train_model(train_h, train_s)\n",
    "mail_result = model.spamprob(focused_mail)\n",
    "print(mail_result)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "len(set(focused_mail))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "s_vocab_arange = range(2000,3000)\n",
    "atk_total = 50"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": false
   },
   "outputs": [],
   "source": [
    "model = train_model(train_h, train_s)\n",
    "\n",
    "def focused_atk():\n",
    "    scores = []\n",
    "    print(\"\\n-------------我是分割线------------\\n\")\n",
    "    focused_chars = list(set(focused_mail))\n",
    "    for n in range(atk_total):\n",
    "        # 创建攻击邮件\n",
    "        h_num = np.random.choice(range(len(focused_chars)))\n",
    "        s_num = np.random.choice(s_vocab_arange)\n",
    "        h_chars = np.random.choice(focused_chars, h_num)\n",
    "        s_chars = np.random.choice(s_vocab_set, s_num)\n",
    "        chars = np.hstack((h_chars, s_chars))\n",
    "        np.random.shuffle(chars)\n",
    "        atk_mail = chars.tolist()\n",
    "        print(\"已创建第 {} 封攻击邮件，邮件长度为 {}\".format(n, len(atk_mail)))\n",
    "        \n",
    "        # 送入邮件防护系统进行识别\n",
    "        mail_result = model.spamprob(atk_mail)\n",
    "        \n",
    "        # 根据得分对模型进行正常的再训练\n",
    "        if mail_result <= HAM_CUTOFF:\n",
    "            print(\"攻击异常\")\n",
    "            model.learn_ham(atk_mail)\n",
    "        elif mail_result >= SPAM_CUTOFF:\n",
    "            print(\"攻击正常\")\n",
    "            model.learn_spam(atk_mail)\n",
    "            \n",
    "        # 结果检测\n",
    "        focused_result = model.spamprob(focused_mail)\n",
    "        scores.append(focused_result)\n",
    "        \n",
    "        if n % 10 == 0:\n",
    "            print(focused_result)\n",
    "            \n",
    "    return scores\n",
    "\n",
    "focused_scores = focused_atk()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "plot_scores(focused_scores)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### 2.4  Pseudospam Attack"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "pseudospam_mail = np.random.choice(test_s)\n",
    "model = train_model(train_h, train_s)\n",
    "mail_result = model.spamprob(pseudospam_mail)\n",
    "print(mail_result)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "len(set(pseudospam_mail))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "h_vocab_arange = range(300,600)\n",
    "atk_total = 50"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "model = train_model(train_h, train_s)\n",
    "\n",
    "def pseudospam_atk():\n",
    "    scores = []\n",
    "    print(\"\\n-------------我是分割线------------\\n\")\n",
    "    pseudospam_chars = list(set(pseudospam_mail))\n",
    "    for n in range(atk_total):\n",
    "        # 创建攻击邮件\n",
    "        h_num = np.random.choice(h_vocab_arange)\n",
    "        s_num = np.random.choice(range(len(pseudospam_chars)))\n",
    "        h_chars = np.random.choice(h_vocab_set, h_num)\n",
    "        s_chars = np.random.choice(pseudospam_chars, s_num)\n",
    "        chars = np.hstack((h_chars, s_chars))\n",
    "        np.random.shuffle(chars)\n",
    "        atk_mail = chars.tolist()\n",
    "        print(\"已创建第 {} 封攻击邮件，邮件长度为 {}\".format(n, len(atk_mail)))\n",
    "        \n",
    "        # 送入邮件防护系统进行识别\n",
    "        mail_result = model.spamprob(atk_mail)\n",
    "        \n",
    "        # 根据得分对模型进行正常的再训练\n",
    "        if mail_result <= HAM_CUTOFF:\n",
    "            print(\"攻击正常\")\n",
    "            model.learn_ham(atk_mail)\n",
    "        elif mail_result >= SPAM_CUTOFF:\n",
    "            print(\"攻击异常\")\n",
    "            model.learn_spam(atk_mail)\n",
    "            \n",
    "        # 结果检测\n",
    "        pseudospam_result = model.spamprob(pseudospam_mail)\n",
    "        scores.append(pseudospam_result)\n",
    "        \n",
    "        if n % 10 == 0:\n",
    "            print(pseudospam_result)\n",
    "            \n",
    "    return scores\n",
    "\n",
    "pseudospam_scores = pseudospam_atk()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "model.spamprob(pseudospam_mail) < 0.2"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "plot_scores(pseudospam_scores)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### 2.5 RONI Defence"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def en_train_model(atk_mail): \n",
    "    # 送入邮件防护系统进行识别\n",
    "    mail_result = model.spamprob(atk_mail)\n",
    "\n",
    "    # 根据得分对模型进行正常的再训练\n",
    "    if mail_result <= HAM_CUTOFF:\n",
    "        model.learn_ham(atk_mail)\n",
    "        score = eval(model, test_h, test_s)\n",
    "        print(\"初始化分数：{} | 当前分数：{}\".format(init_score, score))\n",
    "        if init_score - score > 0.01:\n",
    "            print(\"发现攻击，已防护\")\n",
    "            model.unlearn_ham(atk_mail)\n",
    "    elif mail_result >= SPAM_CUTOFF: \n",
    "        model.learn_spam(atk_mail)\n",
    "        score = eval(model, test_h, test_s)\n",
    "        print(\"初始化分数：{} | 当前分数：{}\".format(init_score, score))\n",
    "        if init_score - score > 0.01:\n",
    "            print(\"发现攻击，已防护\")\n",
    "            model.unlearn_spam(atk_mail)\n",
    "    else:\n",
    "        score = eval(model, test_h, test_s)\n",
    "\n",
    "    return score\n",
    "\n",
    "\n",
    "def dictionary_atk():\n",
    "    init_score = eval(model, test_h, test_s)\n",
    "    h_vocab_arange = range(1000,2000)\n",
    "    s_vocab_arange = range(1000,2000)\n",
    "    atk_total = 2000\n",
    "    scores = []\n",
    "    print(\"\\n-------------我是分割线------------\\n\")\n",
    "    \n",
    "    for n in range(atk_total):\n",
    "        # 创建攻击邮件\n",
    "        h_num = np.random.choice(h_vocab_arange)\n",
    "        s_num = np.random.choice(s_vocab_arange)\n",
    "        h_chars = np.random.choice(h_vocab_set, h_num)\n",
    "        s_chars = np.random.choice(s_vocab_set, s_num)\n",
    "        chars = np.hstack((h_chars, s_chars))\n",
    "        np.random.shuffle(chars)\n",
    "        atk_mail = chars.tolist()\n",
    "        #print(\"已创建第 {} 封攻击邮件，邮件长度为 {}\".format(n, len(atk_mail)))\n",
    "        \n",
    "        # 送入邮件防护系统进行识别\n",
    "        score = en_train_model(atk_mail)\n",
    "        scores.append(score)\n",
    "            \n",
    "    print(\"\\n-------------我是分割线------------\\n\")\n",
    "    print(\"攻击后模型对测试数据识别准确率为：\")\n",
    "    result = eval(model, test_h, test_s)\n",
    "    return scores"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "model = train_model(train_h, train_s)\n",
    "scores = dictionary_atk()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.7.4"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
